import json
from py2neo import Graph, Node, Relationship
from concurrent.futures import ThreadPoolExecutor
import time

# 连接到Neo4j数据库
# graph = Graph("http://192.168.86.128:7474/", user="neo4j", password="12345678", name="neo4j")  # 替换为实际的数据库连接信息
graph = Graph("http://localhost:7474/",auth=("neo4j", "123456"))

# 定义处理单行数据的函数
def process_data(data):
    # 提取src_ip、dst_ip和dst_port字段
    src_ip = data["src_ip"]
    dst_ip = data["dst_ip"]
    dst_port = data["dst_port"]

    # graph = Graph("http://192.168.86.128:7474/", user="neo4j", password="12345678", name="neo4j")

    # 检查关系是否存在
    query = """
    MATCH (start:IP {ip: $src_ip})-[rel:CONNECTS_TO]->(end:IP {ip: $dst_ip})
    RETURN rel
    """
    result = graph.evaluate(query, src_ip=src_ip, dst_ip=dst_ip)

    if result is None:
        # 创建或获取src_ip节点
        src_node = Node("IP", ip=src_ip)
        graph.merge(src_node, "IP", "ip")

        # 创建或获取dst_ip节点
        dst_node = Node("IP", ip=dst_ip)
        graph.merge(dst_node, "IP", "ip")

        # 创建关系及属性
        rel = Relationship(src_node, "CONNECTS_TO", dst_node, times=1, dst_port=dst_port)
        graph.create(rel)
    else:
        # 更新关系的times属性值
        query = """
        MATCH (start:IP {ip: $src_ip})-[rel:CONNECTS_TO]->(end:IP {ip: $dst_ip})
        SET rel.times = rel.times + 1
        """
        graph.run(query, src_ip=src_ip, dst_ip=dst_ip)

t1 = time.time()

# 读取data.json文件
with open("C:/Users/yinwe/Desktop/data.json", "r") as file:
    lines = file.readlines()


for line in lines:
    data = json.loads(line)
    # 使用线程池处理数据
    with ThreadPoolExecutor() as executor:
        executor.submit(process_data, data)
    # process_data(data)

# with ThreadPoolExecutor() as executor:
#     for line in lines:
#         data = json.loads(line)
#         # 使用线程池处理数据
#         executor.submit(process_data, data)
#         # process_data(data)

t2 = time.time()
print('time'+str(t2-t1)+'s')